#include "RSA.h"
#include "zbase64.h"

#include <openssl/rsa.h>
#include <openssl/pem.h>
using namespace dakuang;

// 加解密函数签名
typedef int (*RSA_encryptOrDecrypt)(int flen, const unsigned char *from, unsigned char *to, RSA *rsa, int padding);

// 生成密钥对，输出为PEM字符串
bool CRSA::genKeyStrings(std::string& strPrivateKeyPEMString, std::string& strPublicKeyPEMString)
{
    // 生成RSA
    RSA* rsa = (RSA*)__genKey(1024);
    if (rsa == NULL)
        return false;

    // 输出私钥
    {
        BIO* bio = BIO_new(BIO_s_mem());
        int ret = PEM_write_bio_RSAPrivateKey(bio, rsa, NULL, NULL, 0, NULL, NULL);
        if (ret != 1)
        {
            BIO_free(bio);
            RSA_free(rsa);
            return false;
        }

        char sBuf[1024] = {0};
        int bytes = BIO_read(bio, sBuf, 1024);
        if (bytes <= 0)
        {
            BIO_free(bio);
            RSA_free(rsa);
            return false;
        }

        BIO_free(bio);
        strPrivateKeyPEMString.assign(sBuf, bytes);
    }

    // 输出公钥
    {
        BIO* bio = BIO_new(BIO_s_mem());
        int ret = PEM_write_bio_RSAPublicKey(bio, rsa);
        if (ret != 1)
        {
            BIO_free(bio);
            RSA_free(rsa);
            return false;
        }

        char sBuf[1024] = {0};
        int bytes = BIO_read(bio, sBuf, 1024);
        if (bytes <= 0)
        {
            BIO_free(bio);
            RSA_free(rsa);
            return false;
        }

        BIO_free(bio);
        strPublicKeyPEMString.assign(sBuf, bytes);
    }

    RSA_free(rsa);
    return true;
}

// 使用PEM私钥字符串加密
bool CRSA::encryptByPrivatePEMString(const std::string& strIn, std::string& strOut, const std::string& strKeyPEMString)
{
    // 加载私钥
    BIO* bio = BIO_new_mem_buf((const void*)strKeyPEMString.data(), strKeyPEMString.size());
    RSA* rsa = PEM_read_bio_RSAPrivateKey(bio, NULL, NULL, NULL);
    BIO_free(bio);

    if (rsa == NULL)
        return false;

    // 使用私钥加密
    bool b = __encryptOrDecrypt(strIn, strOut, rsa, (const void*)RSA_private_encrypt, true);
    RSA_free(rsa);
    return b;
}

// 使用PEM私钥字符串解密
bool CRSA::decryptByPrivatePEMString(const std::string& strIn, std::string& strOut, const std::string& strKeyPEMString)
{
    // 加载私钥
    BIO* bio = BIO_new_mem_buf((const void*)strKeyPEMString.data(), strKeyPEMString.size());
    RSA* rsa = PEM_read_bio_RSAPrivateKey(bio, NULL, NULL, NULL);
    BIO_free(bio);

    if (rsa == NULL)
        return false;

    // 检查密文大小是否为分组的整数倍
    int keySize = RSA_size(rsa);
    int blockSize = keySize;
    if ( (strIn.size() % blockSize) != 0 )
        return false;

    // 使用私钥解密
    bool b = __encryptOrDecrypt(strIn, strOut, rsa, (const void*)RSA_private_decrypt, false);
    RSA_free(rsa);
    return b;
}

// 使用PEM公钥字符串加密
bool CRSA::encryptByPublicPEMString(const std::string& strIn, std::string& strOut, const std::string& strKeyPEMString)
{
    // 加载公钥
    BIO* bio = BIO_new_mem_buf((const void*)strKeyPEMString.data(), strKeyPEMString.size());
    RSA* rsa = PEM_read_bio_RSAPublicKey(bio, NULL, NULL, NULL);
    BIO_free(bio);

    if (rsa == NULL)
        return false;

    // 使用公钥加密
    bool b = __encryptOrDecrypt(strIn, strOut, rsa, (const void*)RSA_public_encrypt, true);
    RSA_free(rsa);
    return b;
}

// 使用PEM公钥字符串解密
bool CRSA::decryptByPublicPEMString(const std::string& strIn, std::string& strOut, const std::string& strKeyPEMString)
{
    // 加载公钥
    BIO* bio = BIO_new_mem_buf((const void*)strKeyPEMString.data(), strKeyPEMString.size());
    RSA* rsa = PEM_read_bio_RSAPublicKey(bio, NULL, NULL, NULL);
    BIO_free(bio);

    if (rsa == NULL)
        return false;

    // 检查密文大小是否为分组的整数倍
    int keySize = RSA_size(rsa);
    int blockSize = keySize;
    if ( (strIn.size() % blockSize) != 0 )
        return false;

    // 使用公钥解密
    bool b = __encryptOrDecrypt(strIn, strOut, rsa, (const void*)RSA_public_decrypt, false);
    RSA_free(rsa);
    return b;
}

// 生成RSA密钥结构
void* CRSA::__genKey(int nBits)
{
    RSA* rsa = RSA_new();

    BIGNUM* bne = BN_new();
    BN_set_word(bne, RSA_F4);

    int ret = RSA_generate_key_ex(rsa, nBits, bne, NULL);

    BN_free(bne);

    if (ret != 1)
    {
        RSA_free(rsa);
        return NULL;
    }

    return (void*)rsa;
}

// 使用RSA执行加密或解密
bool CRSA::__encryptOrDecrypt(const std::string& strIn, std::string& strOut, const void* pRSA, const void* pFunc, bool bEncrypt)
{
    const RSA_encryptOrDecrypt encryptOrDecrypt = (const RSA_encryptOrDecrypt)pFunc;

    // 计算加密块大小
    int nKeySize = RSA_size((RSA*)pRSA);
    int nBlockSize = bEncrypt ? (nKeySize - RSA_PKCS1_PADDING_SIZE) : nKeySize;

    const unsigned char* pIn = (const unsigned char*)strIn.data();
    int nInSize = strIn.size();

    unsigned char* pBuf = new unsigned char[nKeySize];

    // 分组迭代加密
    for (int i = 0; i < nInSize; )
    {
        int nBlockLen = (nInSize - i > nBlockSize ? nBlockSize : nInSize - i);

        int ret = encryptOrDecrypt(nBlockLen, pIn + i, pBuf, (RSA*)pRSA, RSA_PKCS1_PADDING);
        if (ret <= 0)
        {
            delete[] pBuf;
            return false;
        }

        strOut.append((char*)pBuf, ret);
        i += nBlockLen;
    }

    delete[] pBuf;
    return true;
}
